package data_store

import (
	"context"

	"go.opentelemetry.io/collector/consumer"
	"go.opentelemetry.io/collector/pdata/pcommon"
	"go.opentelemetry.io/collector/pdata/ptrace"
	semconv "go.opentelemetry.io/otel/semconv/v1.21.0"
)

type DataStore struct {
	traceTable *TraceTable
	spanTable  *SpanTable
	nameTable  *NameTable
}

func NewDataStore(maxStore uint32) *DataStore {
	return &DataStore{
		traceTable: newTraceTable(maxStore),
		spanTable:  newSpanTable(maxStore),
		nameTable:  newNameTable(),
	}
}

func (ds *DataStore) parseServiceInfo(attrMap pcommon.Map) (string, string) {
	serviceName := ""
	serviceVersion := ""

	value, ok := attrMap.Get(string(semconv.ServiceNameKey))
	if ok {
		serviceName = value.AsString()
	}
	value, ok = attrMap.Get(string(semconv.ServiceVersionKey))
	if ok {
		serviceVersion = value.AsString()
	}
	return serviceName, serviceVersion
}

func (ds *DataStore) convertSpan(span ptrace.Span, serviceName, serviceVersion string) *SpanRecord {
	eventList := []*EventRecord{}
	for i := 0; i < span.Events().Len(); i++ {
		ev := span.Events().At(i)
		eventList = append(eventList, &EventRecord{
			Name:    ev.Name(),
			Time:    ev.Timestamp().AsTime().UnixNano() / 1e6,
			AttrMap: ev.Attributes(),
		})
	}
	return &SpanRecord{
		SpanId:         span.SpanID().String(),
		ParentSpanId:   span.ParentSpanID().String(),
		TraceId:        span.TraceID().String(),
		ServiceName:    serviceName,
		ServiceVersion: serviceVersion,
		SpanName:       span.Name(),
		StartTime:      span.StartTimestamp().AsTime().UnixNano() / 1e6,
		EndTime:        span.EndTimestamp().AsTime().UnixNano() / 1e6,
		SpanKind:       span.Kind(),
		AttrMap:        span.Attributes(),
		EventList:      eventList,
	}
}

func (ds *DataStore) insert(resSpan ptrace.ResourceSpans) {
	serviceName, serviceVersion := ds.parseServiceInfo(resSpan.Resource().Attributes())
	for i := 0; i < resSpan.ScopeSpans().Len(); i++ {
		traceIdMap := map[string]bool{}
		traceRootSpanMap := map[string]string{}
		spanList := resSpan.ScopeSpans().At(i)
		for j := 0; j < spanList.Spans().Len(); j++ {
			spanRecord := ds.convertSpan(spanList.Spans().At(j), serviceName, serviceVersion)
			ds.spanTable.Insert(spanRecord)
			traceIdMap[spanRecord.TraceId] = true
			if spanRecord.ParentSpanId == "" {
				traceRootSpanMap[spanRecord.TraceId] = spanRecord.SpanId
				ds.nameTable.Insert(NameItem{
					ServiceName:  serviceName,
					RootSpanName: spanRecord.SpanName,
				})
			}
		}
		for traceId, _ := range traceIdMap {
			oldTraceId := ds.traceTable.RemoveOldTrace()
			if oldTraceId != "" {
				ds.spanTable.RemoveByTrace(oldTraceId)
			}
			ds.traceTable.Insert(&TraceRecord{
				TraceId: traceId,
			})
		}
		for traceId, rootSpanId := range traceRootSpanMap {
			ds.traceTable.UpdateRootSpanId(traceId, rootSpanId)
		}
	}
}

func (ds *DataStore) ConsumeTraces(_ context.Context, td ptrace.Traces) error {
	spans := td.ResourceSpans()
	for i := 0; i < spans.Len(); i++ {
		ds.insert(spans.At(i))
	}
	return nil
}

func (ds *DataStore) Capabilities() consumer.Capabilities {
	return consumer.Capabilities{
		MutatesData: false,
	}
}

func (ds *DataStore) GetTraceTable() *TraceTable {
	return ds.traceTable
}

func (ds *DataStore) GetSpanTable() *SpanTable {
	return ds.spanTable
}

func (ds *DataStore) GetNameTable() *NameTable {
	return ds.nameTable
}
